import os
import yaml
import argparse
from train import run_train
from test_pic import run_test

ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
yaml_path = '{}/codes/config.yaml'.format(ROOT_PATH)


def get_args_from_yaml(yaml_path, mode='pre_train_t_net'):
    yaml_file = open(yaml_path, "r", encoding="utf-8")
    file_data = yaml_file.read()
    yaml_file.close()

    # 加载数据流，返回字典类型数据
    y = yaml.load(file_data, Loader=yaml.FullLoader)
    cfg = y[mode]

    if mode == 'test':
        parser = argparse.ArgumentParser(description='human matting')
        parser.add_argument('--model', default='{}/{}'.format(ROOT_PATH, cfg['model']), help='preTrained model')
        parser.add_argument('--test_pic_path', default='{}/{}/test'.format(ROOT_PATH, cfg['test_pic_path']),
                            help='test picture path')
        parser.add_argument('--output_path', default='{}/{}'.format(ROOT_PATH, cfg['output_path']), help='result path')
        parser.add_argument('--size', type=int, default=cfg['size'], help='input size')

        args = parser.parse_args()
        print(args)
    else:
        parser = argparse.ArgumentParser(description='Fast portrait matting !')
        parser.add_argument('--dataDir', default='{}/{}'.format(ROOT_PATH, cfg['dataDir']), help='dataset directory')
        parser.add_argument('--saveDir', default='{}/{}'.format(ROOT_PATH, cfg['saveDir']), help='model save dir')
        parser.add_argument('--trainData', default=cfg['trainData'], help='train dataset name')

        parser.add_argument('--load', default=cfg['load'], help='save model dir')
        parser.add_argument('--finetuning', action='store_true', default=cfg['finetuning'],
                            help='finetuning the training')

        parser.add_argument('--nThreads', type=int, default=cfg['nThreads'], help='number of threads for data loading')
        parser.add_argument('--train_batch', type=int, default=cfg['train_batch'], help='input batch size for train')
        parser.add_argument('--patch_size', type=int, default=cfg['patch_size'], help='patch size for train')
        parser.add_argument('--lr', type=float, default=cfg['lr'], help='learning rate')
        parser.add_argument('--lrDecay', type=int, default=cfg['lrDecay'])
        parser.add_argument('--lrdecayType', default=cfg['lrdecayType'])
        parser.add_argument('--nEpochs', type=int, default=cfg['nEpochs'], help='number of epochs to train')
        parser.add_argument('--save_epoch', type=int, default=cfg['save_epoch'], help='number of epochs to save model')

        parser.add_argument('--train_phase', default=cfg['train_phase'], help='train phase')

        args = parser.parse_args()
        print(args)
    return args


if __name__ == '__main__':
    run_train(get_args_from_yaml(yaml_path, mode='pre_train_t_net'))
    run_train(get_args_from_yaml(yaml_path, mode='end_to_end'))
    run_test(get_args_from_yaml(yaml_path, mode='test'))
